import torch.nn
import torch.nn.functional as F
from tqdm import tqdm

def evaluate_PSFUnet(net, dataloader, device):
    net.eval()
    num_val_batches = len(dataloader)
    criterion = torch.nn.L1Loss()
    total_loss = 0
    # iterate over the validation set
    with tqdm(dataloader, total=num_val_batches, desc='Validation round', unit='batch', leave=False) as pbar:
        for batch in dataloader:
            label_true = batch['label']
            # move images and labels to correct device and type
            # image = image.to(device=device, dtype=torch.float32)
            label_true = label_true.to(device=device, dtype=torch.float32)

            with torch.no_grad():
                # predict the label
                label_pred = net(label_true)
                loss = criterion(label_true, label_pred)

                # update the progress bar
                pbar.update(label_true.shape[0])
                pbar.set_postfix(**{'loss(batch)': loss.item()})
                total_loss += loss
    net.train()

    # Fixes a potential division by zero error
    if num_val_batches == 0:
        return total_loss
    return total_loss / num_val_batches